#MSS plugin for pymol! Created by Peter Daldrop 13/2/2010
#ver 0.3 all basic features, new: no change in colour by pressing next and back, 
#customizable C-atom colours, variable table width, pdb or mol2 input,columnwise sortability with reverse option,
#coming hopefully soon - planned features:
# customisable labels, mark lines, export lines.

import colorsys,sys,os,string
from pymol import cmd,stored
from Tkinter import *
import tkMessageBox
import tkFileDialog
import tkFont
class MSS:
##############################################    
#      Initialise Application window         #
##############################################
    def __init__(self, master): 
        #creating neccessary variables
        self.listwidth=85
        self.infile='' #file containing spreadsheet
        self.data='' #var to hold the spreadsheet data
        self.table=[] #list to hold data lines as lists
        self.path='' #var containing path to molecule files
        self.molfiles=[] # var containing molecule file names
        self.presfiles=[] #var containing filenames found in path dir
        self.display='lines' #display style var, lines by defauld
        self.molcol=StringVar() #molecule C-atom Color, Cyan by default
        self.molcol.set('cyan') #initialise default
        self.Font=tkFont.Font(family='Monospace',size=12)#to have a font with equal character spacing
        self.suffix='.mol2'#contains the file suffix for the selected database type
        self.colorder=1 #column by which to sort the spreadsheet
        self.entryval=StringVar()#links to entry field for column which is sorted against
        self.inversval=0
        #create Frames
        self.mainfram=Frame(roota, bg='blue')
        self.subfram=Frame(self.mainfram)
        self.sub2=Frame(self.mainfram)
        self.sub3=Frame(self.mainfram)
        self.sub4=Frame(self.mainfram)
        #create Frame content
        self.loader=Button(self.sub3,text='Load',command=self.loadit)#triggers spreadsheet loading dialog
        self.style=Menubutton(self.sub3,text='Style')#to hold menu for selecting Display style
        self.style.grid()
        self.mcol=Menubutton(self.sub3,text='Colour')
        self.mcol.grid()
        self.input=Menubutton(self.sub3,text='Input')
        self.input.grid()
        self.style.styleM=Menu(self.style)#add menu to menubutton
        self.mcol.colourM=Menu(self.mcol)
        self.input.inpM=Menu(self.input)
        self.style["menu"]  =  self.style.styleM
        self.style.styleM.add_command(label='Lines',command=self.setline)#add show as lines option
        self.style.styleM.add_command(label='Sticks',command=self.setstick)#add show as lines option
        self.style.styleM.add_command(label='Surface',command=self.setsurface)#add show as lines option
        self.mcol["menu"] = self.mcol.colourM
        self.mcol.colourM.add_command(label='cyan',command=self.setcyan)
        self.mcol.colourM.add_command(label='blue',command=self.setblue)
        self.mcol.colourM.add_command(label='green',command=self.setgreen)
        self.mcol.colourM.add_command(label='red',command=self.setred)
        self.mcol.colourM.add_command(label='hotpink',command=self.sethotpink)
        self.mcol.colourM.add_command(label='pink',command=self.setpink)
        self.mcol.colourM.add_command(label='white',command=self.setwhite)
        self.input["menu"]=self.input.inpM
        self.input.inpM.add_command(label='mol2',command=self.mol2in)
        self.input.inpM.add_command(label='pdb',command=self.pdbin)
        self.info=Button(self.sub3, text='Info',command=self.Information)#shows some general info about the plugin
        self.sortit=Button(self.sub3,text='Sort', command=self.sortcol)#sorts list by column selected in Entry field
        self.sortlab=Label(self.sub3,text='by col:') #label for Entry field
        self.colnum=Entry(self.sub3,textvariable=self.entryval,width=2) #Entry field for column number which is sorted against
        self.colnum.insert(0,'2') #set default value
        self.entryval.set(self.colnum.get())#set variable to entry value
        self.inverlab=Label(self.sub3,text='reverse: ')
        self.invers=Checkbutton(self.sub3, onvalue=1,offvalue=0,command=self.cb)
        #the label contains the column names as needed for Dock.3.5.54 results
        self.label=Label(self.sub4,text='Name         Rank     Score      E_el        E_vdw      E_des_1    E_des_2',font=self.Font) 
        self.scroller=Scrollbar(self.subfram,orient=VERTICAL)#scroll bar to scroll the list vertically
        self.scroller2=Scrollbar(self.mainfram,orient=HORIZONTAL)
        self.showB=Button(self.sub2,text='Show',command=self.ShowM)#to show selected molecule
        self.forwB=Button(self.sub2,text='Next',command=self.Forw)#to show the next molecule after the selected one
        self.backB=Button(self.sub2,text='Back',command=self.Back)#to show the molecule before the selected one
        #the box containing the data lines
        self.lines=Listbox(self.subfram,selectmode='SINGLE',width=self.listwidth,yscrollcommand=self.scroller.set,xscrollcommand=self.scroller2.set,font=self.Font)
        #packing
        self.mainfram.pack(fill=X,expand=1) #insert mainframe into application window
        self.sub3.pack(side=TOP,fill=X)#insert frame to contain loader and info button
        self.loader.pack(side=LEFT)#insert loader button on left side
        self.input.pack(side=LEFT)
        self.style.pack(side=LEFT)#insert style menu
        self.mcol.pack(side=LEFT) #insert color menu
        self.sortit.pack(side=LEFT)#insert sort button
        self.sortlab.pack(side=LEFT)#insert entry label
        self.colnum.pack(side=LEFT)#insert column sorting entry
        self.inverlab.pack(side=LEFT)
        self.invers.pack(side=LEFT)
        self.info.pack(side=RIGHT)#insert info button on right side
        self.sub2.pack()#insert frame to hold buttons
        self.showB.pack(side=LEFT)#insert show button
        self.backB.pack(side=LEFT)#insert next button
        self.forwB.pack(side=LEFT)#insert back button
        self.sub4.pack(fill=X)
        self.label.pack(side=LEFT)#insert label in mainframe
        self.subfram.pack(fill=X,expand=1)#insert frame to hold listbox and scrollbar
        self.lines.pack(side=LEFT,fill=X,expand=1)#insert listbox
        self.scroller.pack(side=RIGHT, fill=Y)#insert scrollbar
        self.scroller["command"]  = self.lines.yview#set it to scroll the listbox vertically
        self.scroller2.pack(fill=X,expand=1)
        self.scroller2["command"]=self.lines.xview
    
##############################################    
#          Defining Functions                #
##############################################       
    def cb(self):
        if self.inversval==0:
            self.inversval=1
        else:
            self.inversval=0
        #print self.inversval
    #Function to set display style variable to lines (called by menu selection)
    def setline(self):
        self.display='lines'

    #Function to set display style variable to sticks (called by menu selection)
    def setstick(self):
        self.display='sticks'
    #Function to set display style variable to surface (called by menu selection)
    def setsurface(self):
        self.display='surface'  
    #Functions to set the C-atom colour. Triggered by the menu ...
    def setblue(self):
        self.molcol.set('blue') 
    def setcyan(self):
        self.molcol.set('cyan')
    def setred(self):
        self.molcol.set('red')
    def sethotpink(self):
        self.molcol.set('hotpink')
    def setpink(self):
        self.molcol.set('pink')
    def setwhite(self):
        self.molcol.set('white')
    def setgreen(self):
        self.molcol.set('green')
    def mol2in(self):
        self.suffix='.mol2'
    def pdbin(self):
        self.suffix='.pdb'
    def loadit(self): #used to load the spreadsheet file
        inf=tkFileDialog.askopenfilename(title='Select file') #get spreadsheet input file
        try:
             self.infile=open(inf,'r') #open input file
        except TypeError: #occurs if cancel selected
             return
        except IOError: #if file not present or available for reaTCTAGAding
             tkMessageBox.showerror('IOError','Error opening input file!') #in case of trouble... prevents plugin to crash
             return
        self.data=self.infile.readlines()#put all lines in data variable
        for i in self.data:
            self.table.append(string.split(i.strip(),'\t'))
        self.path=string.strip(self.data[0]) #read first line containing the path to the molecules
        self.presfiles=os.listdir(self.path)
        for i in self.data[1:]: #for all lines containing actual molecule data
            buffwidth=len(i)
            if buffwidth>self.listwidth:
                self.listwidth=buffwidth
            name=i.split('\t') #generates a string array of the different columns in the spreadsheet
            name=name[0]+self.suffix#molecule name is stored in 1st column and mol2 suffix added yielding molecule file name
            if self.presfiles.__contains__(name):
                self.molfiles.append(name)#add to list of molecule names
                currline=string.strip(i)#remove newline char at end
                currline=currline.replace('\t','     ')#reformat to generate space (listbox other wise really prints \t)
                self.lines.insert('end',currline)#insert line into listbox
    def ShowM(self): #show selected molecule
         if len(self.lines.curselection())==0: #if no line selected
             return
         index=int(self.lines.curselection()[0])#get index of selected line
         cmd.delete('molecule')
         cmd.load(self.path+self.molfiles[index],'molecule',1)#load and display molecule
                  #cmd.hide('sticks','molecule')
         cmd.hide('surface','molecule')
         cmd.show(self.display,'molecule')
         cmd.color(self.molcol.get(), '(molecule and name C*)')
    def Forw(self): #show next molecule
         if len(self.lines.curselection())==0:#if no line selected
             return
         index=int(self.lines.curselection()[0])#get index of selected line
         if index>=len(self.molfiles)-1:  #if already displaying last molecule in list
             print 'END OF LIST' #just some command line feedback so that you know...
             return False #don't do anything and stop Function
         self.lines.selection_clear(index,index) #unselect current molecule
         self.lines.selection_set(index+1)#select next molecule
         #print self.table[index+1] #test output for sorting function
         cmd.delete('molecule')
         cmd.load(self.path+self.molfiles[index+1],'molecule',1)#display next molecule
         cmd.hide('sticks','molecule')
         cmd.hide('surface','molecule')
         cmd.show(self.display,'molecule')
         cmd.color(self.molcol.get(), '(molecule and name C*)')
    def Back(self):#show molecule before currently selected
        
         if len(self.lines.curselection())==0:#if no line selected
             return
         index=int(self.lines.curselection()[0])#get index of selected line
         if index==0: #if already displaying first molecule in list
             print 'START OF LIST' #just some command line feedback so that you know...
             return #don't do anything and stop Function
         self.lines.selection_clear(index,index)#unselect current molecule
         self.lines.selection_set(index-1)#select molecule before current one
         cmd.delete('molecule')
         cmd.load(self.path+self.molfiles[index-1],'molecule',1)#delete molecule before current one
         cmd.hide('sticks','molecule')
         cmd.hide('surface','molecule')
         cmd.show(self.display,'molecule')
         cmd.color(self.molcol.get(), '(molecule and name C*)')    
    def sortcol(self):
        self.entryval.set(self.colnum.get()) #get columnnumber to sort by from entry field
        sorted=[]#initialise list
        counter=1 #counter variable
        #self.colorder=int(self.entryval.get())
        #print self.entryval.get()
        try:
            self.colorder=int(self.entryval.get())-1 #try to convert string from field to int
        except ValueError: #if no interger in entry field give error message
            print 'Could not convert Column number for sorting to integer! Please make sure you entered a number!'
            tkMessageBox.showerror('Error!','Could not convert column number for sorting to integer!\nPlease make sure you entered a number!')
        while counter < len(self.data): #iterate through molecule containing lines of data
            try: #try and convert sort column content to sort by to number 
                sorter=float(self.table[counter][self.colorder])
            except ValueError: #if no number conversion possible keep as string
                sorter=self.table[counter][self.colorder]                
            linex=[sorter,string.strip(self.data[counter])] #generate list for this molecule to be sorted by
            sorted.append(linex)       #append molecule to sort list
            counter +=1 #increment counter variable
        sorted.sort() #sort list of molecules by first element in list (column value to be sorted by)
        if self.inversval==1: #if reverse order requested
            sorted.reverse() #reverse list
        counter =0 #reset counter 
        while counter<len(sorted): #iterate through all entries of sorted list
            sorted[counter]=sorted[counter][1]# remove sorting column from entries keeping just the data
            counter +=1 #increment counter
        self.molfiles=[] #reinitialise file list
        self.lines.delete(0,10000000) #delete all lines from list (should be less than 10 million lines in list)
        for i in sorted: #for all lines containing actual molecule data
            #buffwidth=len(i)
            #if buffwidth>self.listwidth:
            #    self.listwidth=buffwidth
            name=i.split('\t') #generates a string array of the different columns in the spreadsheet
            name=name[0]+self.suffix#molecule name is stored in 1st column and mol2 suffix added yielding molecule file name
            if self.presfiles.__contains__(name):
                self.molfiles.append(name)#add to list of molecule names
                currline=string.strip(i)#remove newline char at end
                currline=currline.replace('\t','     ')#reformat to generate space (listbox other wise really prints \t)
                self.lines.insert('end',currline)#insert line into listbox
    def Information(self): #displays some info in a message box
         tkMessageBox.showinfo("MSS","Molecular Spreadsheet for PyMol\n Created by Peter Daldrop\n Version 0.12\n p.daldrop@dundee.ac.uk \n\n\nSpreadsheet format: first line contains full path to .mol2 molecule files, which are named molname.mol2, from second line on tab-separated comlumns. First column must contain molecule name identical as filename of molecule without .mol2 suffix.! All other columns are free for the users choice though built in labels correspond to DOCK3.5.54 docking scores.")
def __init__(self): #initialise Plugin in Pymol
     #add MSS  option in Plugin menu
    self.menuBar.addmenuitem('Plugin', 'command','MSS',label = 'MSS',command = lambda s=self : open_MSS())

##############################################    
#          Main application                  #
############################################## 
def __init__(self):
    self.menuBar.addmenuitem('Plugin', 'command','MSS',label = 'MSS',command = lambda s=self : open_MSS())
def open_MSS(): #upon plugin selection in pymol
    #initialize global Tk root widget
    global roota #create global objext
    roota = Tk() 
    roota.title('MSS by Peter Daldrop')
    #instanciate pyM Class object with root widget as master
    global mss
    mss = MSS(roota)